package com.weixing.mall.core.interceptor;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.util.concurrent.RateLimiter;
import com.weixing.mall.base.wrapper.Result;
import com.weixing.mall.base.wrapper.ResultUtil;
import com.weixing.mall.core.annotation.RequestLimiter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.servlet.HandlerInterceptor;

import javax.annotation.Resource;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.io.OutputStream;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 请求限流拦截器
 *
 * @author xhq
 * @version 1.0
 * @date 2019/10/22 16:46
 */
@Slf4j
@Component
public class RequestLimiterInterceptor implements HandlerInterceptor {

    @Resource
    private ObjectMapper objectMapper;

    /**
     * 不同的方法存放不同的令牌桶
     */
    private final Map<String, RateLimiter> rateLimiterMap = new ConcurrentHashMap<>();

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) {
        try {
            if (handler instanceof HandlerMethod) {
                HandlerMethod handlerMethod = (HandlerMethod) handler;
                RequestLimiter rateLimit = handlerMethod.getMethodAnnotation(RequestLimiter.class);
                //判断是否有注解
                if (rateLimit != null) {
                    // 获取请求url
                    String url = request.getRequestURI();
                    RateLimiter rateLimiter;
                    // 判断map集合中是否有创建好的令牌桶
                    if (!rateLimiterMap.containsKey(url)) {
                        // 创建令牌桶,以n r/s往桶中放入令牌
                        rateLimiter = RateLimiter.create(rateLimit.QPS());
                        rateLimiterMap.put(url, rateLimiter);
                    }
                    rateLimiter = rateLimiterMap.get(url);
                    // 获取令牌
                    boolean acquire = rateLimiter.tryAcquire(rateLimit.timeout(), rateLimit.timeunit());
                    if (acquire) {
                        //获取令牌成功
                        return true;
                    } else {
                        log.warn("请求被限流,url:{}", request.getServletPath());
                        sendResponse(response,ResultUtil.error(rateLimit.msg()));
                        return false;
                    }
                }
            }
            return true;
        } catch (Exception var6) {
           log.warn("限流发生异常",var6);
            sendResponse(response,ResultUtil.error("对不起,请求似乎出现了一些问题,请您稍后重试！"));
            return false;
        }
    }


    private void sendResponse(HttpServletResponse response, Result result) {
        response.setCharacterEncoding("utf-8");
        response.setContentType("text/html;charset=UTF-8");
        try (OutputStream out = response.getOutputStream();
             PrintWriter pw = new PrintWriter(new OutputStreamWriter(out, "utf-8"))) {
            // 返回json格式的提示
            pw.println(objectMapper.writeValueAsString(result));
        } catch (IOException ex) {
            log.warn("", ex);
        }
    }

}